#!/usr/bin/env python3
"""
G3数据集快速开始示例
"""

import pandas as pd
import json

def quick_start():
    """快速开始示例"""

    print("=== G3数据集快速开始 ===\n")

    # 1. 加载Parquet文件
    print("1. 加载数据...")
    train_df = pd.read_parquet("./g3_hf_dataset_parquet/train.parquet")
    print(f"   训练集: {len(train_df):,} 条数据")

    # 2. 查看数据样例
    print("\n2. 数据样例:")
    sample = train_df.iloc[0]
    print(f"   查询ID: {sample['query_id']}")
    print(f"   查询内容: {sample['query'][:100]}...")
    print(f"   可用API数量: {sample['num_available_apis']}")
    print(f"   相关API数量: {sample['num_relevant_apis']}")

    # 3. 获取API信息
    print("\n3. API信息:")
    api_names = sample['api_names']
    print(f"   前5个API: {api_names[:5]}")

    # 4. 获取原始API数据
    print("\n4. 原始API数据 (api_list_ori):")
    original_apis = json.loads(sample['api_list_ori'])
    first_api = original_apis[0]
    print(f"   第一个API: {first_api['api_name']}")
    print(f"   工具名称: {first_api['tool_name']}")
    print(f"   类别: {first_api['category_name']}")
    print(f"   描述: {first_api['api_description'][:100]}...")

    # 5. 获取相关API
    print("\n5. 相关API:")
    relevant_apis = sample['relevant_apis']
    for i, api in enumerate(relevant_apis[:3]):
        print(f"   {i+1}. {api[0]} - {api[1]}")

    # 6. 数据集统计
    print("\n6. 数据集统计:")
    print(f"   总查询数: {len(train_df):,}")
    print(f"   平均API数量: {train_df['num_available_apis'].mean():.2f}")
    print(f"   平均相关API数量: {train_df['num_relevant_apis'].mean():.2f}")
    print(f"   平均查询长度: {train_df['query'].str.len().mean():.2f} 字符")

    # 7. 使用示例代码
    print("\n7. 基本使用代码:")
    print("""
# 导入必要的库
import pandas as pd
import json

# 加载数据
train_df = pd.read_parquet('train.parquet')
val_df = pd.read_parquet('validation.parquet')
test_df = pd.read_parquet('test.parquet')

# 获取单条数据
sample = train_df.iloc[0]
query = sample['query']
api_names = sample['api_names']
relevant_apis = sample['relevant_apis']

# 获取原始API详细信息
original_apis = json.loads(sample['api_list_ori'])

# 批量处理
for idx, row in train_df.head(10).iterrows():
    print(f"查询 {row['query_id']}: {row['query'][:50]}...")
    print(f"相关API: {len(row['relevant_apis'])} 个")
""")

    print("\n=== 完成! ===")

if __name__ == "__main__":
    quick_start()